module CodeGen
  ( genModule
  , encodeProc
  ) where

import Prelude

import AST as A
import Control.Monad.ST (ST)
import Data.Array.ST (STArray, new, push, run)
import Data.ArrayBuffer.Types (Uint8Array)
import Data.Foldable (traverse_)
import Data.FoldableWithIndex (class FoldableWithIndex, foldrWithIndex, traverseWithIndex_)
import Data.Maybe (Maybe(..))
import Wasm.Encode (encodeModule)
import Wasm.Syntax as W

encodeProc :: A.Proc -> Uint8Array
encodeProc = encodeModule <<< genModule

genModule :: A.Proc -> W.Module
genModule proc = W.emptyModule
  { memories = [ { type: { min: 0, max: Nothing } } ]
  , imports =
      [ { module: "js", name: "putchar", desc: W.ImportFunc 0 }
      , { module: "js", name: "getchar", desc: W.ImportFunc 1 }
      , { module: "js", name: "mem", desc: W.ImportMemory { min: 0, max: Nothing } }
      ]
  , exports = [ { name: "run", desc: W.ExportFunc 2 } ]
  , funcs = [ { type: 0, locals: [], body: runGen (_ `genProc` proc) } ]
  , types =
      [ funcType [ W.NumType W.I32 ] []
      , funcType [] [ W.NumType W.I32 ]
      ]
  }

funcType :: Array W.ValType -> Array W.ValType -> W.RecType
funcType arguments results =
  [ { final: true, supertypes: [], ty: W.CompFunc { arguments, results } } ]

runGen :: (forall a. STArray a W.Instruction -> ST a Unit) -> W.Expr
runGen gen = run do
  arr <- new
  gen arr
  pure arr

genProc :: forall a. STArray a W.Instruction -> A.Proc -> ST a Unit
genProc a proc = traverse_ (genInstr a) proc

genInstr :: forall a. STArray a W.Instruction -> A.Instr -> ST a Unit
genInstr a (A.Move i) = do
  pointer a
  p a $ W.I32Const i
  p a $ W.I32Add
  p a $ W.LocalSet 0
genInstr a (A.In i) = do
  pointer a
  p a $ W.Call 1
  store a i
genInstr a (A.Out i) = do
  load a i
  p a $ W.Call 0
genInstr a' (A.Loop { idx, proc }) =
  p a' $ W.Block (W.BlockValType Nothing) $
    [ W.Loop (W.BlockValType Nothing) $ runGen \a -> do
        load a idx
        p a $ W.I32Eqz
        p a $ W.Br_if 1
        genProc a proc
        p a $ W.Br 0
    ]
genInstr a (A.Update m) = do
  traverseWithIndex_ (\i o -> pointer a *> genOper a i o) m
  traverserIndex_ (store a) m

genOper :: forall a. STArray a W.Instruction -> Int -> A.Oper -> ST a Unit
genOper a _ (A.Set e) = genExpr a e
genOper a i (A.Add x) = genExpr a $ A.EAdd (A.ECst x) (A.EDyn i)

genExpr :: forall a. STArray a W.Instruction -> A.Expr -> ST a Unit
genExpr a (A.ECst x) = p a $ W.I32Const x
genExpr a (A.EDyn i) = load a i
genExpr a (A.ENeg e) = genExpr a $ A.ESub (A.ECst 0) e
genExpr a (A.EAdd x y) = do
  genExpr a x
  genExpr a y
  p a $ W.I32Add
genExpr a (A.ESub x y) = do
  genExpr a x
  genExpr a y
  p a $ W.I32Sub
genExpr a (A.EMul x y) = do
  genExpr a x
  genExpr a y
  p a $ W.I32Mul

store :: forall a. STArray a W.Instruction -> Int -> ST a Unit
store a i = p a $ W.I32Store8 { offset: i, align: 0 }

load :: forall a. STArray a W.Instruction -> Int -> ST a Unit
load a i = do
  pointer a
  p a $ W.I32Load8_u { offset: i, align: 0 }

pointer :: forall a. STArray a W.Instruction -> ST a Unit
pointer a = p a $ W.LocalGet 0

p :: forall a v. STArray a v -> v -> ST a Unit
p a v = void $ push v a

traverserIndex_ :: forall i f a b m. FoldableWithIndex i f => Applicative m => (i -> m b) -> f a -> m Unit
traverserIndex_ f = foldrWithIndex (\i _ a -> a <* f i) (pure unit)
